D:\a\csshw\csshw\src\lib.rs
Line | Count | Source |
1 | | //! Cluster SSH tool for Windows inspired by csshX |
2 | | |
3 | | #![deny(clippy::implicit_return)] |
4 | | #![allow(clippy::needless_return, clippy::doc_overindented_list_items)] |
5 | | #![warn(missing_docs)] |
6 | | #![doc(html_no_source)] |
7 | | #![cfg_attr(coverage_nightly, feature(coverage_attribute))] |
8 | | |
9 | | use std::fs::{create_dir, File}; |
10 | | use std::mem; |
11 | | |
12 | | use log::warn; |
13 | | use registry::{value, Data, Hive, Security}; |
14 | | use simplelog::{format_description, ConfigBuilder, LevelFilter, WriteLogger}; |
15 | | use windows::core::PWSTR; |
16 | | use windows::Win32::Foundation::HWND; |
17 | | use windows::Win32::System::Threading::{PROCESS_INFORMATION, STARTUPINFOW}; |
18 | | |
19 | | #[cfg(test)] |
20 | | use mockall::automock; |
21 | | |
22 | | pub mod cli; |
23 | | pub mod client; |
24 | | pub mod daemon; |
25 | | pub mod serde; |
26 | | pub mod utils; |
27 | | |
28 | | use utils::windows::WindowsApi; |
29 | | |
30 | | /// CLSID identifying `conhost.exe` in the registry. |
31 | | /// |
32 | | /// As used in Windows Terminal: |
33 | | /// <https://github.com/microsoft/terminal/blob/v1.22.3232.0/src/propslib/DelegationConfig.hpp#L105> |
34 | | const CLSID_CONHOST: &str = "{B23D10C0-E52E-411E-9D5B-C09FDF709C7D}"; |
35 | | /// CLSID identifying the default configuration in the registry. |
36 | | /// |
37 | | /// The default configuration is "let windows choose". |
38 | | /// Also defined in Windows Terminal: |
39 | | /// <https://github.com/microsoft/terminal/blob/v1.22.3232.0/src/propslib/DelegationConfig.hpp#L104> |
40 | | const CLSID_DEFAULT: &str = "{00000000-0000-0000-0000-000000000000}"; |
41 | | /// Registry path where `DelegationConsole` and `DelegationTerminal` registry keys are stored. |
42 | | /// |
43 | | /// These registry keys store the configuration value for the default terminal application. |
44 | | const DEFAULT_TERMINAL_APP_REGISTRY_PATH: &str = r"Console\%%Startup"; |
45 | | /// `DelegationConsole` registry key. |
46 | | /// |
47 | | /// As used in Windows Terminal: |
48 | | /// <https://github.com/microsoft/terminal/blob/v1.22.3232.0/src/propslib/DelegationConfig.cpp#L29> |
49 | | const DELEGATION_CONSOLE: &str = "DelegationConsole"; |
50 | | /// `DelegationTerminal` registry key. |
51 | | /// |
52 | | /// As used in Windows Terminal: |
53 | | /// <https://github.com/microsoft/terminal/blob/v1.22.3232.0/src/propslib/DelegationConfig.cpp#L30> |
54 | | const DELEGATION_TERMINAL: &str = "DelegationTerminal"; |
55 | | |
56 | | /// Trait for registry operations to enable mocking in tests |
57 | | #[cfg_attr(test, automock)] |
58 | | pub trait Registry { |
59 | | /// Get a string value from the registry |
60 | | fn get_registry_string_value(&self, path: &str, name: &str) -> Option<String>; |
61 | | /// Set a string value in the registry |
62 | | fn set_registry_string_value(&self, path: &str, name: &str, value: &str) -> bool; |
63 | | } |
64 | | |
65 | | /// Default implementation of Registry trait that performs actual Windows registry API calls |
66 | | pub struct DefaultRegistry; |
67 | | |
68 | | #[cfg_attr(coverage_nightly, coverage(off))] |
69 | | impl Registry for DefaultRegistry { |
70 | | fn get_registry_string_value(&self, path: &str, name: &str) -> Option<String> { |
71 | | let key = Hive::CurrentUser |
72 | | .open(path, Security::Read | Security::Write) |
73 | | .ok()?; |
74 | | match key.value(name) { |
75 | | Ok(Data::String(value)) => return Some(value.to_string_lossy()), |
76 | | Ok(_) => panic!("Expected string data for {name} registry value"), |
77 | | Err(value::Error::NotFound(_, _)) => return Some(CLSID_DEFAULT.to_owned()), |
78 | | Err(err) => { |
79 | | warn!("Failed to read {} value from registry: {}", name, err); |
80 | | return None; |
81 | | } |
82 | | } |
83 | | } |
84 | | |
85 | | fn set_registry_string_value(&self, path: &str, name: &str, value: &str) -> bool { |
86 | | if let Ok(key) = Hive::CurrentUser.open(path, Security::Read | Security::Write) { |
87 | | match key.set_value::<String>( |
88 | | name.to_owned(), |
89 | | &Data::String(value.to_owned().try_into().unwrap()), |
90 | | ) { |
91 | | Ok(()) => return true, |
92 | | Err(_) => { |
93 | | warn!("Failed to set registry value {} to {}", name, value); |
94 | | return false; |
95 | | } |
96 | | } |
97 | | } else { |
98 | | return false; |
99 | | } |
100 | | } |
101 | | } |
102 | | |
103 | | /// Return the Window Handle [HWND] for the foreground window associated with the given `process_id`. |
104 | | /// |
105 | | /// If multiple foreground windows are associated with the given `process_id` it is undefined which [HWND] gets returned. |
106 | | /// |
107 | | /// # Arguments |
108 | | /// |
109 | | /// * `windows_api` - Windows API operations implementation |
110 | | /// * `process_id` - ID of the process for which to retrieve the window handle. |
111 | | /// |
112 | | /// # Returns |
113 | | /// |
114 | | /// The Window Handle [HWND] for the window associated with the given `process_id`. |
115 | 5 | pub fn get_console_window_handle<W: WindowsApi>(windows_api: &W, process_id: u32) -> HWND { |
116 | 5 | return windows_api.get_window_handle_for_process(process_id); |
117 | 5 | } |
118 | | |
119 | | /// Create process with command line using the provided API (testable version) |
120 | | /// |
121 | | /// # Arguments |
122 | | /// |
123 | | /// * `api` - Windows API operations implementation |
124 | | /// * `application` - Application name including file extension |
125 | | /// * `command_line` - UTF-16 encoded command line |
126 | | /// |
127 | | /// # Returns |
128 | | /// |
129 | | /// [PROCESS_INFORMATION] of the spawned process or None if failed |
130 | 3 | pub fn create_process<W: WindowsApi>( |
131 | 3 | api: &W, |
132 | 3 | application: &str, |
133 | 3 | command_line: &[u16], |
134 | 3 | ) -> Option<PROCESS_INFORMATION> { |
135 | 3 | let mut startupinfo = STARTUPINFOW { |
136 | 3 | cb: mem::size_of::<STARTUPINFOW>() as u32, |
137 | 3 | ..Default::default() |
138 | 3 | }; |
139 | 3 | let mut process_information = PROCESS_INFORMATION::default(); |
140 | 3 | let mut cmd_line = command_line.to_vec(); |
141 | 3 | let command_line_ptr = PWSTR(cmd_line.as_mut_ptr()); |
142 | | |
143 | 3 | match api.create_process_raw( |
144 | 3 | application, |
145 | 3 | command_line_ptr, |
146 | 3 | &mut startupinfo, |
147 | 3 | &mut process_information, |
148 | 3 | ) { |
149 | 2 | Ok(()) => return Some(process_information), |
150 | 1 | Err(_) => return None, |
151 | | } |
152 | 3 | } |
153 | | |
154 | | /// Trait for file system operations to enable mocking in tests |
155 | | #[cfg_attr(test, automock)] |
156 | | pub trait FileSystem { |
157 | | /// Create a directory |
158 | | fn create_directory(&self, path: &str) -> bool; |
159 | | /// Create a log file |
160 | | fn create_log_file(&self, filename: &str) -> bool; |
161 | | } |
162 | | |
163 | | /// Default implementation of FileSystem trait that performs actual file system operations |
164 | | pub struct ProductionFileSystem; |
165 | | |
166 | | #[cfg_attr(coverage_nightly, coverage(off))] |
167 | | impl FileSystem for ProductionFileSystem { |
168 | | fn create_directory(&self, path: &str) -> bool { |
169 | | return create_dir(path).is_ok() || std::path::Path::new(path).exists(); |
170 | | } |
171 | | |
172 | | fn create_log_file(&self, filename: &str) -> bool { |
173 | | return File::create(filename).is_ok(); |
174 | | } |
175 | | } |
176 | | |
177 | | /// Guard storing previous/old `DelegationConsole` and `DelegationTerminal` registry values. |
178 | | /// |
179 | | /// Configures `conhost.exe` as the default terminal application |
180 | | /// and reverts to the original configuration when being dropped. |
181 | | pub struct WindowsSettingsDefaultTerminalApplicationGuard<R: Registry> { |
182 | | /// Old `DelegationConsole` registry value |
183 | | old_windows_terminal_console: Option<String>, |
184 | | /// Old `DelegationTerminal` registry value |
185 | | old_windows_terminal_terminal: Option<String>, |
186 | | /// Registry operations trait |
187 | | registry: R, |
188 | | } |
189 | | |
190 | | impl<R: Registry> WindowsSettingsDefaultTerminalApplicationGuard<R> { |
191 | | /// Create a new guard with the given registry operations |
192 | | /// |
193 | | /// # Arguments |
194 | | /// |
195 | | /// * `registry` - Registry operations implementation |
196 | | /// |
197 | | /// # Returns |
198 | | /// |
199 | | /// A new guard that will restore registry values on drop |
200 | 11 | pub fn new_with_registry(registry: R) -> Self { |
201 | 11 | let mut guard = WindowsSettingsDefaultTerminalApplicationGuard { |
202 | 11 | old_windows_terminal_console: None, |
203 | 11 | old_windows_terminal_terminal: None, |
204 | 11 | registry, |
205 | 11 | }; |
206 | | |
207 | 3 | if let (Some(console_val), Some(terminal_val)) = ( |
208 | 11 | guard |
209 | 11 | .registry |
210 | 11 | .get_registry_string_value(DEFAULT_TERMINAL_APP_REGISTRY_PATH, DELEGATION_CONSOLE), |
211 | 11 | guard |
212 | 11 | .registry |
213 | 11 | .get_registry_string_value(DEFAULT_TERMINAL_APP_REGISTRY_PATH, DELEGATION_TERMINAL), |
214 | | ) { |
215 | | // No need to change if already set to conhost |
216 | 3 | if console_val == CLSID_CONHOST && terminal_val == CLSID_CONHOST1 { |
217 | 1 | return guard; |
218 | 2 | } |
219 | | |
220 | | // Store old values and set new ones |
221 | 2 | guard.old_windows_terminal_console = Some(console_val); |
222 | 2 | guard.old_windows_terminal_terminal = Some(terminal_val); |
223 | | |
224 | 2 | guard.registry.set_registry_string_value( |
225 | 2 | DEFAULT_TERMINAL_APP_REGISTRY_PATH, |
226 | 2 | DELEGATION_CONSOLE, |
227 | 2 | CLSID_CONHOST, |
228 | | ); |
229 | 2 | guard.registry.set_registry_string_value( |
230 | 2 | DEFAULT_TERMINAL_APP_REGISTRY_PATH, |
231 | 2 | DELEGATION_TERMINAL, |
232 | 2 | CLSID_CONHOST, |
233 | | ); |
234 | | } else { |
235 | 8 | warn!( |
236 | | "Failed to read registry key {}, \ |
237 | | cannot make sure conhost.exe is the configured default terminal application", |
238 | | DEFAULT_TERMINAL_APP_REGISTRY_PATH, |
239 | | ); |
240 | | } |
241 | | |
242 | 10 | return guard; |
243 | 11 | } |
244 | | } |
245 | | |
246 | | impl WindowsSettingsDefaultTerminalApplicationGuard<DefaultRegistry> { |
247 | | /// Create a new guard with production registry operations |
248 | 6 | pub fn new() -> Self { |
249 | 6 | return Self::new_with_registry(DefaultRegistry); |
250 | 6 | } |
251 | | } |
252 | | |
253 | | impl<R: Registry> Default for WindowsSettingsDefaultTerminalApplicationGuard<R> |
254 | | where |
255 | | R: Default, |
256 | | { |
257 | 0 | fn default() -> Self { |
258 | 0 | return Self::new_with_registry(R::default()); |
259 | 0 | } |
260 | | } |
261 | | |
262 | | impl Default for DefaultRegistry { |
263 | 0 | fn default() -> Self { |
264 | 0 | return DefaultRegistry; |
265 | 0 | } |
266 | | } |
267 | | |
268 | | impl<R: Registry> Drop for WindowsSettingsDefaultTerminalApplicationGuard<R> { |
269 | | /// Restore the original default terminal application setting to the registry. |
270 | | /// |
271 | | /// If old values weren't stored, nothing is done. |
272 | 11 | fn drop(&mut self) { |
273 | 2 | if let (Some(old_console), Some(old_terminal)) = ( |
274 | 11 | &self.old_windows_terminal_console, |
275 | 11 | &self.old_windows_terminal_terminal, |
276 | 2 | ) { |
277 | 2 | self.registry.set_registry_string_value( |
278 | 2 | DEFAULT_TERMINAL_APP_REGISTRY_PATH, |
279 | 2 | DELEGATION_CONSOLE, |
280 | 2 | old_console, |
281 | 2 | ); |
282 | 2 | self.registry.set_registry_string_value( |
283 | 2 | DEFAULT_TERMINAL_APP_REGISTRY_PATH, |
284 | 2 | DELEGATION_TERMINAL, |
285 | 2 | old_terminal, |
286 | 2 | ); |
287 | 9 | } |
288 | 11 | } |
289 | | } |
290 | | |
291 | | /// Launch the given console application with the given arguments as a new detached process with its own console window. |
292 | | /// |
293 | | /// Input/Output handles are not being inherited. |
294 | | /// Whichever default terminal application is configured in the windows system settings will be used |
295 | | /// to host the application (i.e. create the window). |
296 | | /// |
297 | | /// # Arguments |
298 | | /// |
299 | | /// * `api` - Windows API implementation |
300 | | /// * `application` - Application name including file extension (`.exe`). |
301 | | /// If the application is not in the `PATH` environment variable, the full path |
302 | | /// must be specified. |
303 | | /// * `args` - List of arguments to the application. |
304 | | /// |
305 | | /// # Returns |
306 | | /// |
307 | | /// [PROCESS_INFORMATION] of the spawned process. |
308 | 10 | pub fn spawn_console_process<W: WindowsApi>( |
309 | 10 | api: &W, |
310 | 10 | application: &str, |
311 | 10 | args: Vec<String>, |
312 | 10 | ) -> Option<PROCESS_INFORMATION> { |
313 | 10 | return api.create_process_with_args(application, args); |
314 | 10 | } |
315 | | |
316 | | /// Initialize the logger. |
317 | | /// |
318 | | /// Makes sure a `logs` directory exists in the current working directory. |
319 | | /// Log filename format: `<utc-time-of-executable-start>_<name>.log`. |
320 | | /// Configures [log_panics]. |
321 | | /// |
322 | | /// # Arguments |
323 | | /// |
324 | | /// * `name` - Will be part of the log filename. |
325 | 0 | pub fn init_logger(name: &str) { |
326 | 0 | init_logger_with_fs(&ProductionFileSystem, name); |
327 | 0 | } |
328 | | |
329 | | /// Initialize the logger with the provided file system operations. |
330 | | /// |
331 | | /// # Arguments |
332 | | /// |
333 | | /// * `fs` - File system operations implementation |
334 | | /// * `name` - Will be part of the log filename |
335 | 9 | pub fn init_logger_with_fs<F: FileSystem>(fs: &F, name: &str) { |
336 | 9 | let utc_now = chrono::offset::Utc::now() |
337 | 9 | .format("%Y-%m-%d_%H-%M-%S.%f") |
338 | 9 | .to_string(); |
339 | | |
340 | 9 | fs.create_directory("logs"); |
341 | | |
342 | 9 | let filename = format!("logs/{utc_now}_{name}.log"); |
343 | 9 | if fs.create_log_file(&filename) { |
344 | 7 | if let Ok(file0 ) = File::create(&filename) { |
345 | 0 | let _ = WriteLogger::init( |
346 | 0 | LevelFilter::Debug, |
347 | 0 | ConfigBuilder::new() |
348 | 0 | .set_time_format_custom(format_description!( |
349 | 0 | "[hour]:[minute]:[second].[subsecond]" |
350 | 0 | )) |
351 | 0 | .build(), |
352 | 0 | file, |
353 | 0 | ); |
354 | 0 | log_panics::init(); |
355 | 7 | } |
356 | 2 | } |
357 | 9 | } |
358 | | |
359 | | /// Detect if application was launched from Windows Explorer (GUI) vs command line using the provided console API. |
360 | | /// |
361 | | /// Returns true if launched from GUI (separate console), false if from existing console. |
362 | | /// Based on: <https://devblogs.microsoft.com/oldnewthing/20160125-00/?p=92922> |
363 | | /// |
364 | | /// # Arguments |
365 | | /// |
366 | | /// * `windows_api` - Windows API operations implementation |
367 | | /// |
368 | | /// # Returns |
369 | | /// |
370 | | /// * `true` - Application was launched from GUI (Explorer, double-click, etc.) |
371 | | /// * `false` - Application was launched from existing console (command line) |
372 | 12 | pub fn is_launched_from_gui<W: WindowsApi>(windows_api: &W) -> bool { |
373 | 12 | return windows_api.get_console_attached_process_count() == 1; |
374 | 12 | } |
375 | | |
376 | | #[cfg(test)] |
377 | | #[path = "./tests/test_lib.rs"] |
378 | | mod test_lib; |